// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#ifdef POLLUX_ENABLE_HDFS
#include "melon/concurrency/ConcurrentHashMap.h"

#include <pollux/common/config/Config.h>
#include <pollux/connectors/hive/storage_adapters/hdfs/HdfsFileSystem.h> // @manual
#include <pollux/connectors/hive/storage_adapters/hdfs/HdfsUtil.h> // @manual
#include <pollux/dwio/common/FileSink.h>
#endif

namespace kumo::pollux::filesystems {

#ifdef POLLUX_ENABLE_HDFS
std::mutex mtx;

std::function<std::shared_ptr<
    FileSystem>(std::shared_ptr<const config::ConfigBase>, std::string_view)>
hdfsFileSystemGenerator() {
  static auto filesystemGenerator = [](std::shared_ptr<const config::ConfigBase>
                                           properties,
                                       std::string_view filePath) {
    static melon::ConcurrentHashMap<std::string, std::shared_ptr<FileSystem>>
        filesystems;
    static melon::
        ConcurrentHashMap<std::string, std::shared_ptr<melon::once_flag>>
            hdfsInitiationFlags;
    HdfsServiceEndpoint endpoint =
        HdfsFileSystem::getServiceEndpoint(filePath, properties.get());
    std::string hdfsIdentity = endpoint.identity();
    if (filesystems.find(hdfsIdentity) != filesystems.end()) {
      return filesystems[hdfsIdentity];
    }
    std::unique_lock<std::mutex> lk(mtx, std::defer_lock);
    /// If the init flag for a given hdfs identity is not found,
    /// create one for init use. It's a singleton.
    if (hdfsInitiationFlags.find(hdfsIdentity) == hdfsInitiationFlags.end()) {
      lk.lock();
      if (hdfsInitiationFlags.find(hdfsIdentity) == hdfsInitiationFlags.end()) {
        std::shared_ptr<melon::once_flag> initiationFlagPtr =
            std::make_shared<melon::once_flag>();
        hdfsInitiationFlags.insert(hdfsIdentity, initiationFlagPtr);
      }
      lk.unlock();
    }
    melon::call_once(
        *hdfsInitiationFlags[hdfsIdentity].get(),
        [&properties, endpoint, hdfsIdentity]() {
          auto filesystem =
              std::make_shared<HdfsFileSystem>(properties, endpoint);
          filesystems.insert(hdfsIdentity, filesystem);
        });
    return filesystems[hdfsIdentity];
  };
  return filesystemGenerator;
}

std::function<std::unique_ptr<pollux::dwio::common::FileSink>(
    const std::string&,
    const pollux::dwio::common::FileSink::Options& options)>
hdfsWriteFileSinkGenerator() {
  static auto hdfsWriteFileSink =
      [](const std::string& fileURI,
         const pollux::dwio::common::FileSink::Options& options) {
        if (HdfsFileSystem::isHdfsFile(fileURI)) {
          std::string pathSuffix =
              getHdfsPath(fileURI, HdfsFileSystem::kScheme);
          auto fileSystem =
              filesystems::getFileSystem(fileURI, options.connectorProperties);
          return std::make_unique<dwio::common::WriteFileSink>(
              fileSystem->openFileForWrite(pathSuffix),
              fileURI,
              options.metricLogger,
              options.stats);
        }
        return static_cast<std::unique_ptr<dwio::common::WriteFileSink>>(
            nullptr);
      };

  return hdfsWriteFileSink;
}
#endif

void registerHdfsFileSystem() {
#ifdef POLLUX_ENABLE_HDFS
  registerFileSystem(HdfsFileSystem::isHdfsFile, hdfsFileSystemGenerator());
  dwio::common::FileSink::registerFactory(hdfsWriteFileSinkGenerator());
#endif
}

} // namespace kumo::pollux::filesystems
